"""
Enhanced Image Encoder模块
功能：将消息嵌入图像中，集成PatchGAN对抗训练提升隐蔽性
这是FakeTagger系统的第2个模块
"""

import torch
import torch.nn as nn
import torch.nn.functional as F
from encoder import ResNetUNet  # 使用原有的encoder
from patch_discriminator import PatchDiscriminator, RaLSGANLoss

class EnhancedImageEncoder(nn.Module):
    """
    增强版图像编码器
    在原有ResNet-UNet编码器基础上(集成PatchGAN判别器进行对抗训练)
    这是来自SepMark项目的创新点
    """
    
    def __init__(self, 
                 message_size=30, 
                 use_discriminator=True, 
                 discriminator_weight=0.1, 
                 device='cuda', 
                 lr=0.0002, 
                 beta1=0.5):
        """
        初始化增强版图像编码器
        
        Args:
            message_size: 消息长度（冗余消息长度）
            use_discriminator: 是否使用判别器进行对抗训练
            discriminator_weight: 对抗损失权重
            device: 设备类型
            lr: 学习率
            beta1: Adam优化器参数
        """
        super(EnhancedImageEncoder, self).__init__()
        
        self.message_size = message_size
        self.use_discriminator = use_discriminator
        self.discriminator_weight = discriminator_weight
        self.device = device
        
        # 原有的ResNet-UNet编码器（保持原有架构）
        self.encoder = ResNetUNet(message_size)
        
        print(f"Enhanced Image Encoder initialized:")
        print(f"  - Message size: {message_size}")
        print(f"  - Use discriminator: {use_discriminator}")
        print(f"  - Discriminator weight: {discriminator_weight}")
        
        # PatchGAN判别器（来自SepMark的创新）
        if use_discriminator:
            self.discriminator = PatchDiscriminator(input_nc=3, ndf=64, n_layers=3)
            
            # 损失函数
            self.criterion_adversarial = RaLSGANLoss()  # 相对平均最小二乘GAN损失
            self.criterion_reconstruction = nn.L1Loss()
            
            # 优化器（分别为编码器和判别器）
            self.opt_encoder = torch.optim.Adam(
                self.encoder.parameters(), 
                lr=lr, betas=(beta1, 0.999)
            )
            self.opt_discriminator = torch.optim.Adam(
                self.discriminator.parameters(), 
                lr=lr, betas=(beta1, 0.999)
            )
            
            print(f"  - PatchGAN discriminator initialized")
            print(f"  - Using RaLSGAN loss function")
        else:
            self.discriminator = None
            self.opt_encoder = torch.optim.Adam(
                self.encoder.parameters(), 
                lr=lr, betas=(beta1, 0.999)
            )
            print(f"  - Standard training mode (no discriminator)")
        
        # 移动到设备
        self.device_obj = torch.device(device if torch.cuda.is_available() else 'cpu')
        print(f"Enhanced Image Encoder device: {self.device_obj}")
        
        # 移动编码器
        self.encoder = self.encoder.to(self.device_obj)
        print(f"Encoder moved to: {next(self.encoder.parameters()).device}")
        
        # 移动判别器（如果有）
        if self.discriminator is not None:
            self.discriminator = self.discriminator.to(self.device_obj)
            print(f"Discriminator moved to: {next(self.discriminator.parameters()).device}")
        
        # 验证设备正确性
        encoder_device = next(self.encoder.parameters()).device
        if str(encoder_device) != str(self.device_obj):
            print(f"⚠️ Warning: Encoder device mismatch! Expected: {self.device_obj}, Got: {encoder_device}")
        else:
            print(f"✅ Encoder correctly on device: {encoder_device}")
            
        if self.discriminator is not None:
            disc_device = next(self.discriminator.parameters()).device
            if str(disc_device) != str(self.device_obj):
                print(f"⚠️ Warning: Discriminator device mismatch! Expected: {self.device_obj}, Got: {disc_device}")
            else:
                print(f"✅ Discriminator correctly on device: {disc_device}")
    
    def to(self, device):
        """移动模块到指定设备"""
        super().to(device)
        
        # 更新设备属性
        self.device_obj = device
        
        # 移动编码器
        self.encoder = self.encoder.to(device)
        
        # 移动判别器
        if self.discriminator is not None:
            self.discriminator = self.discriminator.to(device)
        
        return self
    
    def encode(self, images, messages):
        """
        编码：在图像中嵌入水印消息
        
        Args:
            images: 输入图像 (batch_size, 3, H, W)
            messages: 要嵌入的消息 (batch_size, message_size)
            
        Returns:
            编码后的图像 (batch_size, 3, H, W)
        """
        # 确保输入在正确的设备上
        if hasattr(self, 'device_obj'):
            images = images.to(self.device_obj)
            messages = messages.to(self.device_obj)
        
        return self.encoder(images, messages)
    
    def standard_train_step(self, images, messages):
        """
        标准训练步骤（不使用对抗训练）
        
        Args:
            images: 输入图像
            messages: 消息
            
        Returns:
            dict: 损失信息
            torch.Tensor: 编码后的图像
        """
        # 前向传播
        encoded_images = self.encode(images, messages)
        
        # 重建损失
        recon_loss = self.criterion_reconstruction(encoded_images, images)
        
        # 反向传播
        self.opt_encoder.zero_grad()
        recon_loss.backward()
        self.opt_encoder.step()
        
        losses = {
            'encoder_recon_loss': recon_loss.item(),
            'total_encoder_loss': recon_loss.item()
        }
        
        return losses, encoded_images
    
    def adversarial_train_step(self, images, messages):
        """
        对抗训练步骤（使用PatchGAN判别器）
        
        Args:
            images: 输入图像 (batch_size, 3, H, W)
            messages: 消息 (batch_size, message_size)
            
        Returns:
            dict: 包含各种损失的字典
            torch.Tensor: 编码后的图像
        """
        if not self.use_discriminator:
            return self.standard_train_step(images, messages)
        
        # 编码图像
        encoded_images = self.encode(images, messages)
        
        # ========================
        # 训练判别器 (Discriminator)
        # ========================
        # 启用判别器参数的梯度计算
        for p in self.discriminator.parameters():
            p.requires_grad = True
            
        self.opt_discriminator.zero_grad()
        
        # 判别器对真实图像和编码图像的预测
        d_real = self.discriminator(images)
        d_fake = self.discriminator(encoded_images.detach())  # detach防止梯度回传到编码器
        
        # 使用RaLSGAN损失
        d_loss = self.criterion_adversarial(d_real, d_fake, is_discriminator=True)
        
        # 判别器反向传播
        d_loss.backward()
        self.opt_discriminator.step()
        
        # ========================  
        # 训练编码器 (Generator/Encoder)
        # ========================
        # 禁用判别器参数的梯度计算
        for p in self.discriminator.parameters():
            p.requires_grad = False
            
        self.opt_encoder.zero_grad()
        
        # 编码器的对抗损失
        g_real = self.discriminator(images)
        g_fake = self.discriminator(encoded_images)
        
        # 对抗损失：让编码器生成的图像能够欺骗判别器
        g_adv_loss = self.criterion_adversarial(g_real, g_fake, is_discriminator=False)
        
        # 重建损失：编码后的图像应该与原图相似
        g_recon_loss = self.criterion_reconstruction(encoded_images, images)
        
        # 总编码器损失
        g_total_loss = self.discriminator_weight * g_adv_loss + g_recon_loss
        
        # 编码器反向传播
        g_total_loss.backward()
        self.opt_encoder.step()
        
        # 整理损失信息
        losses = {
            'discriminator_loss': d_loss.item(),
            'encoder_adv_loss': g_adv_loss.item(),
            'encoder_recon_loss': g_recon_loss.item(),
            'total_encoder_loss': g_total_loss.item()
        }
        
        return losses, encoded_images
    
    def train_step(self, images, messages):
        """
        统一的训练步骤接口
        
        Args:
            images: 输入图像
            messages: 消息
            
        Returns:
            dict: 损失信息
            torch.Tensor: 编码后的图像
        """
        if self.use_discriminator:
            return self.adversarial_train_step(images, messages)
        else:
            return self.standard_train_step(images, messages)
    
    def save_encoder_state(self, path):
        """保存编码器状态"""
        state_dict = {
            'encoder': self.encoder.state_dict(),
            'message_size': self.message_size,
            'use_discriminator': self.use_discriminator,
            'discriminator_weight': self.discriminator_weight
        }
        
        if self.use_discriminator:
            state_dict.update({
                'discriminator': self.discriminator.state_dict(),
                'opt_encoder': self.opt_encoder.state_dict(),
                'opt_discriminator': self.opt_discriminator.state_dict()
            })
        else:
            state_dict['opt_encoder'] = self.opt_encoder.state_dict()
        
        torch.save(state_dict, path)
        print(f"Enhanced Image Encoder saved to {path}")
    
    def load_encoder_state(self, path):
        """加载编码器状态"""
        state_dict = torch.load(path, map_location=self.device)
        
        self.encoder.load_state_dict(state_dict['encoder'])
        
        if self.use_discriminator and 'discriminator' in state_dict:
            self.discriminator.load_state_dict(state_dict['discriminator'])
            if 'opt_encoder' in state_dict:
                self.opt_encoder.load_state_dict(state_dict['opt_encoder'])
            if 'opt_discriminator' in state_dict:
                self.opt_discriminator.load_state_dict(state_dict['opt_discriminator'])
        elif 'opt_encoder' in state_dict:
            self.opt_encoder.load_state_dict(state_dict['opt_encoder'])
        
        print(f"Enhanced Image Encoder loaded from {path}")

def test_enhanced_image_encoder():
    """测试增强版图像编码器"""
    print("=== Testing Enhanced Image Encoder ===")
    
    device = 'cuda' if torch.cuda.is_available() else 'cpu'
    
    # 测试参数
    batch_size = 4
    message_size = 60  # 冗余消息长度（30*2）
    image_size = 256
    
    # 创建测试数据
    test_images = torch.randn(batch_size, 3, image_size, image_size)
    test_messages = torch.randn(batch_size, message_size)
    
    if device == 'cuda':
        test_images = test_images.cuda()
        test_messages = test_messages.cuda()
    
    print(f"Test data - Images: {test_images.shape}, Messages: {test_messages.shape}")
    
    # 测试标准编码器
    print("\n--- Testing Standard Encoder ---")
    standard_encoder = EnhancedImageEncoder(
        message_size=message_size, 
        use_discriminator=False, 
        device=device
    )
    
    standard_encoder.train()
    std_losses, std_encoded = standard_encoder.train_step(test_images, test_messages)
    
    print("Standard encoder losses:")
    for key, value in std_losses.items():
        print(f"  {key}: {value:.6f}")
    print(f"Encoded images shape: {std_encoded.shape}")
    
    # 测试对抗编码器
    print("\n--- Testing Adversarial Encoder ---")
    adv_encoder = EnhancedImageEncoder(
        message_size=message_size, 
        use_discriminator=True, 
        discriminator_weight=0.1,
        device=device
    )
    
    adv_encoder.train()
    adv_losses, adv_encoded = adv_encoder.train_step(test_images, test_messages)
    
    print("Adversarial encoder losses:")
    for key, value in adv_losses.items():
        print(f"  {key}: {value:.6f}")
    print(f"Encoded images shape: {adv_encoded.shape}")
    
    # 测试推理模式
    print("\n--- Testing Inference Mode ---")
    adv_encoder.eval()
    with torch.no_grad():
        inference_encoded = adv_encoder.encode(test_images, test_messages)
    print(f"Inference encoded shape: {inference_encoded.shape}")
    
    print("Enhanced Image Encoder test completed!\n")

if __name__ == "__main__":
    test_enhanced_image_encoder()